package rolling

import (
	"sync"
	"time"
)

type Number struct {
	rwLock   *sync.RWMutex
	buckets  map[int64]*NumberBucket
	interval time.Duration
	keepTime time.Duration
}

type NumberBucket struct {
	Value int64
}

func NewNumber(interval, keepTime time.Duration) *Number {
	if interval == 0 {
		interval = 100 * time.Millisecond
	}
	if keepTime == 0 {
		keepTime = 10 * time.Second
	}
	return &Number{
		rwLock:   new(sync.RWMutex),
		buckets:  make(map[int64]*NumberBucket),
		interval: interval,
		keepTime: keepTime,
	}
}

func (n *NumberBucket) Increment(num int64) {
	n.Value += num
}

// 按100ms设置一个bucket
func (n *Number) getCurrentKey() int64 {
	return time.Now().UnixNano() / n.interval.Nanoseconds()
}

func (n *Number) getMinKey(keepTime time.Duration) int64 {
	return n.getCurrentKey() - keepTime.Nanoseconds()/n.interval.Nanoseconds()
}

func (n *Number) getCurrentBucket() *NumberBucket {
	currentKey := n.getCurrentKey()
	if bucket, ok := n.buckets[currentKey]; ok {
		return bucket
	} else {
		n.buckets[currentKey] = &NumberBucket{}
		return n.buckets[currentKey]
	}
}

func (n *Number) removeOldBuckets() {
	minKey := n.getMinKey(n.keepTime)
	for key := range n.buckets {
		if key < minKey {
			delete(n.buckets, key)
		}
	}
}

func (n *Number) Increment(num int64) {
	n.rwLock.Lock()
	defer n.rwLock.Unlock()
	currentBucket := n.getCurrentBucket()
	currentBucket.Increment(num)
	n.removeOldBuckets()
}

func (n *Number) Sum() int64 {

	var sum int64 = 0
	minKey := n.getMinKey(n.keepTime)

	n.rwLock.RLock()
	defer n.rwLock.RUnlock()

	for key, bucket := range n.buckets {
		if key >= minKey {
			sum += bucket.Value
		}
	}
	return sum
}

func (n *Number) SumLastSecond() int64 {
	var sum int64 = 0
	minKey := n.getMinKey(time.Second)

	n.rwLock.RLock()
	defer n.rwLock.RUnlock()

	for key, bucket := range n.buckets {
		if key >= minKey {
			sum += bucket.Value
		}
	}
	return sum
}
